'''
    按顺序返回进程池的子进程运行结果
'''

from multiprocessing import Pool # JoinableQueue, Process, Queue 
import multiprocessing
import sys
import time
import traceback
import threading


class QueueHandler:
    def __init__(self, queue):
        self.queue = queue
    
    def write(self, *args):
        for message in args:
            self.queue.put({"header":"message", "data":message})
    
    def flush(self):  # the flush is necessary，though in most case it was unuseful
        pass


def my_fun(num1, num2, num3, a=None, b=None):
    print("num", num1, num2, num3)
    print("char", a, b)
    

def my_process(*args, **kwds):
    fun = kwds.pop('fun')
    q = kwds.pop("q")
    sys.stdout = QueueHandler(q) # 重定向输出，数据会自动加上换行符
    try:
        result = fun(*args, **kwds)
        q.put({"header":"result", "data":result}) 
    except KeyboardInterrupt:
        print("+++++++++++++++++++程序被打断++++++++++++++++++++++++++")
        q.put({"header":"killed", "data":""})
        return
    except Exception as e:
        q.put({"header":"error", "data":str(e)+traceback.format_exc()})
    q.join()
    

class MyPool:
    def __init__(self) -> None:
        '''
            to implement a function that running a muilti process which can output log in sequence
            1. the first process print first, then one by one (first-in first-out)
            2. the first process result return first, then one by one
        '''
        self.pool = Pool(4)
        self.p_list = []
        self.q_list = []
        self.result = []
        self.task_done = []
        self.start = None
        self.stop = False
        self._terminate = False
        self.print = print
    
    def terminate(self):
        self._terminate = True
        self.pool.terminate()
        for p in self.pool._pool:  # 注意：这是一个内部属性，不推荐使用，仅为示例说明。
            p.kill()     
        self.print("Pool terminated.")
    
    def apply_async(self, fun, *args, **kwargs):
        '''
            fun: the function of muilti process to run
            args and kwargs: the argument of the fun
        '''
        kwargs['fun'] =  fun
        q = multiprocessing.Manager().Queue()
        kwargs['q'] = q
        p = self.pool.apply_async(my_process, args=args, kwds=kwargs)
        self.p_list.append(p)
        self.q_list.append(q)
        self.result.append([])
        self.task_done.append(False)
        if not self.start:
            self.start = threading.Thread(target=self.communicate)
            self.start.start()
            
    def read(self):
        self.stop = True
        if self.start is not None:
            self.start.join()
        ret,err = [],[]
        for i in self.result:
            if i[-1]['header'] == 'result':
                ret.append(i[-1]['data'])
            elif i[-1]['header'] == 'error':
                err.append(i[-1]['data'])
        return ret, err
        
    def communicate(self):
        '''
            ensure to apply all the processes first
        '''
        cu_idx = 0
        while True:    
            for i,p in enumerate(self.p_list):
                q = self.q_list[i]
                if i>=cu_idx and not q.empty():
                    data = q.get()
                    q.task_done()
                    self.result[i].append(data)
                    # print current process message 
                    if cu_idx==i and data['header']=='message':
                        self.print(data['data'], end='')    
                    if data["header"]=='result' or data["header"]=='error':
                        self.task_done[i] = True 
                    if data["header"]=='killed':
                        self.terminate()
                if p.ready():
                    self.task_done[i] = True   
                    
                if self.task_done[i] and cu_idx==i: # cu_idx may larger than len(self.p_list) because user didn't add new p currently
                    cu_idx += 1
                    # print the message that already saved in previous
                    if len(self.result)>cu_idx:
                        for data in self.result[cu_idx]:
                            if data['header']=='message':
                                
                                self.print(data['data'], end='')                        
            if (self.stop and cu_idx==len(self.p_list)) or self._terminate:
                break
            time.sleep(0.001)
        # print("close pool")
        self.pool.close()
        self.pool.join()
    
        
if __name__=="__main__":
    my_pool = MyPool()
    my_pool.apply_async(my_fun, 1,2,3, a='a', b='b')
    result = my_pool.read()
    # 用法参考ODBCReport的check.py脚本
    